{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU",
  "gpuClass": "standard"
 },
 "cells": [
  {
   "cell_type": "code",
   "source": [
    "%%bash\n",
    "pip install transformers -q\n",
    "pip install sentencepiece -q\n",
    "pip install peft -q\n",
    "pip install accelerate -q\n",
    "pip install bitsandbytes -q"
   ],
   "metadata": {
    "id": "AiUynN6DqHsk"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "import transformers\n",
    "from peft import PeftModel\n",
    "\n",
    "model_name = \"reasonwang/flan-t5-xl-8bit\"\n",
    "peft_model_id = \"reasonwang/flan-alpaca-gpt4-lora-xl\"\n",
    "\n",
    "tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)\n",
    "base_model = transformers.AutoModelForSeq2SeqLM.from_pretrained(model_name)\n",
    "peft_model = PeftModel.from_pretrained(base_model, peft_model_id)"
   ],
   "metadata": {
    "id": "gGXQ2YmgkUic"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "prompt = \"If you are the president of a developing country, what you will do to make your country better?\"\n",
    "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
    "for k, v in inputs.items():\n",
    "    inputs[k] = v.to(\"cuda\")\n",
    "outputs = peft_model.generate(**inputs, max_length=256, do_sample=True)\n",
    "text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]\n",
    "print(text)"
   ],
   "metadata": {
    "id": "ypY5w_1m1_0d",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  }
 ]
}
